import argparse

parser = argparse.ArgumentParser()

parser.add_argument('--trainer', type=str, default='commonTrainer', help='trainer name')
parser.add_argument('--net_name', type=str, default='hairnet', help='net name')
parser.add_argument('--input_shape', type=int, nargs="+", default=(224, 224, 3), help='input image shape')
parser.add_argument('--output_shape', type=int, nargs="+", default=(224, 224, 3), help='input image shape')


parser.add_argument('--num_classes', type=int, default=1, help='number of model output channels')
parser.add_argument('-bs', '--batch_size', type=int, default=16, help='batch size')
parser.add_argument('--test_batch_size', type=int, default=8, help='batch size')

parser.add_argument('--max_epoch', type=int, default=100, help='number of epochs to train for')
parser.add_argument('-lr', '--learning_rate', type=float, default=0.1, help='learning rate')

parser.add_argument('--sample_intervel', type=int, default=200, help='step of saving sample images')
parser.add_argument('--log_intervel', type=int, default=20, help='step of saving sample images')

parser.add_argument('--checkpoint_root_dir', default='snapshot', help="path to saved models (to continue training)")
parser.add_argument('--model_path', default='', help="path to saved models (to continue training)")

parser.add_argument('--sample_root_dir', default='train_samples', help='folder to save predict image when training')
parser.add_argument('--num_workers', type=int, default=4, help='number of data loading workers')


parser.add_argument('--trainBasedir', default='/home/hanson/work/HairMattingTools/train', help='dataset directory')
parser.add_argument('--trainList', default='/home/hanson/work/HairMattingTools/trainlist.txt', help='train img list')

parser.add_argument('--testBasedir', default='/home/hanson/work/HairMattingTools/train', help='dataset directory')
parser.add_argument('--testList', default='/home/hanson/work/HairMattingTools/testlist.txt', help='test img list')

parser.add_argument('-rm', '--return_mask', action="store_true", default=False, help='return mask from dataset')
parser.add_argument('-rb', '--return_blurmask', action="store_true", default=False, help='return blurmask from dataset')
parser.add_argument('-rt', '--return_trimap', action="store_true", default=False, help='return trimap from dataset')
parser.add_argument('-ra', '--return_alpha', action="store_true", default=False, help='return alpha from dataset')



parser.add_argument('--refine', action="store_true", default=False, help='finetune')

parser.add_argument('--reset_train_sample', action="store_true", default=False, help='reset train sample dir')

def get_config():
    return parser.parse_args()
